"""
麻将数据生成器
用于生成训练和测试数据
"""
import os
import json
import random
from typing import Dict, List, Optional, Tuple, Union
import logging
import numpy as np
import datetime
from tqdm import tqdm

from sichuanmajiang.engine.cards import TileSet, Tile
from sichuanmajiang.engine.rules import can_pong, can_kong, can_chi, can_win
from sichuanmajiang.engine.game import MahjongGame, Player
from sichuanmajiang.data.data_collector import DataCollector


class MahjongDataGenerator:
    """
    麻将数据生成器
    """
    def __init__(self,
                 output_dir: str,
                 num_games: int = 1000,
                 game_rules: Optional[Dict] = None,
                 seed: Optional[int] = None):
        """
        Args:
            output_dir: 输出目录
            num_games: 生成的游戏数量
            game_rules: 游戏规则配置
            seed: 随机种子
        """
        self.output_dir = output_dir
        self.num_games = num_games
        self.game_rules = game_rules or {}
        
        # 设置随机种子
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)
        
        # 初始化日志
        self.logger = logging.getLogger('MahjongDataGenerator')
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 数据收集器
        self.collector = DataCollector()
    
    def generate_game_data(self) -> List[Dict]:
        """
        生成单局游戏数据
        
        Returns:
            游戏数据列表
        """
        # 创建游戏实例
        game = MahjongGame(rules=self.game_rules)
        
        # 进行游戏
        game_data = []
        while not game.is_game_over():
            # 获取当前游戏状态
            current_state = game.get_state()
            
            # 获取当前玩家
            current_player = game.current_player
            
            # 选择动作 - 使用简单的随机策略
            possible_actions = game.get_possible_actions()
            if possible_actions:
                action = random.choice(possible_actions)
                action_name = action[0] if isinstance(action, tuple) else action
                
                # 记录状态和动作
                state_data = {
                    'state': current_state,
                    'player': current_player.name,
                    'hand_tiles': [str(tile) for tile in current_player.hand_tiles],
                    'action': action_name,
                    'timestamp': datetime.datetime.now().isoformat()
                }
                
                # 如果是打牌动作，记录打出的牌
                if action_name == 'discard' and len(action) > 1:
                    discarded_tile = action[1]
                    state_data['discarded_tile'] = str(discarded_tile)
                
                game_data.append(state_data)
                
                # 执行动作
                game.execute_action(action)
            else:
                self.logger.warning(f"无可执行动作，游戏结束")
                break
        
        # 记录游戏结果
        result = game.get_result()
        for record in game_data:
            record['game_result'] = result
        
        return game_data
    
    def generate_expert_game(self) -> List[Dict]:
        """
        生成专家水平的游戏数据（基于规则的策略）
        
        Returns:
            游戏数据列表
        """
        # 创建游戏实例
        game = MahjongGame(rules=self.game_rules)
        
        # 使用规则型AI
        for i in range(4):
            game.players[i] = RuleBasedAIPlayer(f"AI_{i}")
        
        # 进行游戏
        game_data = []
        while not game.is_game_over():
            # 获取当前游戏状态
            current_state = game.get_state()
            
            # 获取当前玩家
            current_player = game.current_player
            
            # 选择动作 - 使用基于规则的策略
            action = current_player.choose_action(game)
            action_name = action[0] if isinstance(action, tuple) else action
            
            # 记录状态和动作
            state_data = {
                'state': current_state,
                'player': current_player.name,
                'hand_tiles': [str(tile) for tile in current_player.hand_tiles],
                'action': action_name,
                'timestamp': datetime.datetime.now().isoformat(),
                'strategy': 'rule_based'
            }
            
            # 如果是打牌动作，记录打出的牌
            if action_name == 'discard' and len(action) > 1:
                discarded_tile = action[1]
                state_data['discarded_tile'] = str(discarded_tile)
            
            game_data.append(state_data)
            
            # 执行动作
            game.execute_action(action)
        
        # 记录游戏结果
        result = game.get_result()
        for record in game_data:
            record['game_result'] = result
        
        return game_data
    
    def generate_dataset(self, expert_only: bool = False) -> List[Dict]:
        """
        生成完整的数据集
        
        Args:
            expert_only: 是否只生成专家水平的数据
            
        Returns:
            数据集列表
        """
        dataset = []
        
        # 进度条
        pbar = tqdm(total=self.num_games, desc="生成游戏数据")
        
        for i in range(self.num_games):
            try:
                if expert_only:
                    game_data = self.generate_expert_game()
                else:
                    # 80%专家数据，20%随机数据
                    if random.random() < 0.8:
                        game_data = self.generate_expert_game()
                    else:
                        game_data = self.generate_game_data()
                
                # 添加到数据集
                dataset.extend(game_data)
                
                # 更新进度
                pbar.update(1)
                
                # 每100局保存一次中间结果
                if (i + 1) % 100 == 0:
                    temp_file = os.path.join(self.output_dir, f"temp_dataset_{i+1}.json")
                    self._save_dataset(temp_file, dataset)
                    self.logger.info(f"已生成 {i+1} 局游戏，临时保存到 {temp_file}")
                    
            except Exception as e:
                self.logger.error(f"生成第 {i+1} 局游戏时出错: {e}")
        
        pbar.close()
        
        return dataset
    
    def _save_dataset(self, file_path: str, dataset: List[Dict]) -> None:
        """
        保存数据集到文件
        
        Args:
            file_path: 文件路径
            dataset: 数据集
        """
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(dataset, f, ensure_ascii=False, indent=2)
    
    def run(self, 
            expert_only: bool = False,
            split_train_val: bool = True,
            val_ratio: float = 0.2) -> Tuple[str, Optional[str]]:
        """
        运行数据生成器
        
        Args:
            expert_only: 是否只生成专家水平的数据
            split_train_val: 是否分割训练集和验证集
            val_ratio: 验证集比例
            
        Returns:
            (训练集文件路径, 验证集文件路径)
        """
        self.logger.info(f"开始生成 {self.num_games} 局麻将数据...")
        
        # 生成数据集
        dataset = self.generate_dataset(expert_only=expert_only)
        
        self.logger.info(f"数据生成完成，共 {len(dataset)} 条记录")
        
        # 文件名
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        
        if split_train_val and val_ratio > 0:
            # 分割数据集
            random.shuffle(dataset)
            val_size = int(len(dataset) * val_ratio)
            
            val_dataset = dataset[:val_size]
            train_dataset = dataset[val_size:]
            
            # 保存文件
            train_file = os.path.join(self.output_dir, f"train_data_{timestamp}.json")
            val_file = os.path.join(self.output_dir, f"val_data_{timestamp}.json")
            
            self._save_dataset(train_file, train_dataset)
            self._save_dataset(val_file, val_dataset)
            
            self.logger.info(f"训练集保存到 {train_file}, 共 {len(train_dataset)} 条记录")
            self.logger.info(f"验证集保存到 {val_file}, 共 {len(val_dataset)} 条记录")
            
            return train_file, val_file
        else:
            # 保存完整数据集
            full_file = os.path.join(self.output_dir, f"full_data_{timestamp}.json")
            self._save_dataset(full_file, dataset)
            
            self.logger.info(f"完整数据集保存到 {full_file}, 共 {len(dataset)} 条记录")
            
            return full_file, None


class RuleBasedAIPlayer(Player):
    """
    基于规则的AI玩家
    """
    def __init__(self, name: str):
        """
        Args:
            name: 玩家名称
        """
        super().__init__(name)
    
    def choose_action(self, game: MahjongGame) -> Tuple[str, ...]:
        """
        选择动作
        
        Args:
            game: 游戏实例
            
        Returns:
            动作元组
        """
        # 获取可能的动作
        possible_actions = game.get_possible_actions()
        
        if not possible_actions:
            return ('pass',)
        
        # 1. 优先选择胡牌
        for action in possible_actions:
            if action[0] == 'win':
                return action
        
        # 2. 其次选择杠牌
        for action in possible_actions:
            if action[0] == 'kong':
                return action
        
        # 3. 再次选择碰牌
        for action in possible_actions:
            if action[0] == 'pong':
                return action
        
        # 4. 然后选择吃牌
        for action in possible_actions:
            if action[0] == 'chi':
                # 选择最优的吃牌方式
                return self._select_best_chi(action)
        
        # 5. 最后选择打牌
        for action in possible_actions:
            if action[0] == 'discard':
                # 使用贪婪策略选择要打出的牌
                return self._select_best_discard(game)
        
        # 默认动作
        return possible_actions[0]
    
    def _select_best_chi(self, chi_actions: List[Tuple]) -> Tuple:
        """
        选择最优的吃牌方式
        
        Args:
            chi_actions: 吃牌动作列表
            
        Returns:
            最优的吃牌动作
        """
        # 简单策略：优先选择能够形成更多顺子的组合
        if len(chi_actions) == 1:
            return chi_actions[0]
        
        # 随机选择
        return random.choice(chi_actions)
    
    def _select_best_discard(self, game: MahjongGame) -> Tuple[str, Tile]:
        """
        选择最优的打出牌
        
        Args:
            game: 游戏实例
            
        Returns:
            打牌动作
        """
        # 1. 统计每张牌的数量
        tile_counts = {}
        for tile in self.hand_tiles:
            if tile not in tile_counts:
                tile_counts[tile] = 0
            tile_counts[tile] += 1
        
        # 2. 分类牌：孤张、对子、刻子、顺子候选
        isolated_tiles = []  # 孤张（只出现一次）
        pairs = []  # 对子（出现两次）
        triplets = []  # 刻子（出现三次）
        
        for tile, count in tile_counts.items():
            if count == 1:
                isolated_tiles.append(tile)
            elif count == 2:
                pairs.append(tile)
            elif count == 3:
                triplets.append(tile)
        
        # 3. 优先打出孤张
        if isolated_tiles:
            # 对孤张进行排序，优先打出价值最低的牌
            # 简单策略：考虑附近是否有牌，形成顺子
            self.hand_tiles.sort()  # 先排序便于分析
            
            # 计算每张孤张的价值（简单计算附近是否有牌）
            tile_values = {}
            for tile in isolated_tiles:
                value = self._calculate_tile_value(tile, self.hand_tiles)
                tile_values[tile] = value
            
            # 选择价值最低的牌（优先打出）
            worst_tile = min(tile_values.items(), key=lambda x: x[1])[0]
            return ('discard', worst_tile)
        
        # 4. 如果没有孤张，考虑拆对子（如果对子太多）
        if pairs and len(pairs) > 2:
            # 选择价值最低的对子拆
            pair_values = {}
            for tile in pairs:
                value = self._calculate_tile_value(tile, self.hand_tiles)
                pair_values[tile] = value
            
            worst_tile = min(pair_values.items(), key=lambda x: x[1])[0]
            return ('discard', worst_tile)
        
        # 5. 如果没有更好的选择，随机打一张牌
        if self.hand_tiles:
            return ('discard', self.hand_tiles[0])
        
        # 保险起见
        return ('pass',)
    
    def _calculate_tile_value(self, tile: Tile, hand_tiles: List[Tile]) -> float:
        """
        计算牌的价值
        
        Args:
            tile: 要评估的牌
            hand_tiles: 手牌
            
        Returns:
            牌的价值（越高越好）
        """
        value = 0
        
        # 计算附近是否有牌（形成顺子的可能性）
        # 所有牌都是序数牌
        # 检查前后各两张牌
        for offset in [-2, -1, 1, 2]:
            neighbor_number = tile.value + offset
            if 1 <= neighbor_number <= 9:
                # 创建相邻牌并检查是否在手牌中
                neighbor_tile = Tile(tile.suit, neighbor_number)
                if neighbor_tile in hand_tiles:
                    # 相邻牌的价值递减
                    value += 1.0 / abs(offset)
        
        return value


class DataAugmenter:
    """
    数据增强器
    """
    @staticmethod
    def augment_data(data: List[Dict], 
                    augment_factor: int = 2) -> List[Dict]:
        """
        对数据进行增强
        
        Args:
            data: 原始数据
            augment_factor: 增强倍数
            
        Returns:
            增强后的数据
        """
        augmented_data = []
        
        for _ in range(augment_factor):
            for record in data:
                # 创建副本
                augmented = record.copy()
                
                # 应用不同的增强策略
                if 'hand_tiles' in augmented:
                    # 1. 牌的顺序打乱（不影响语义）
                    if random.random() < 0.5:
                        hand_tiles = augmented['hand_tiles'].copy()
                        random.shuffle(hand_tiles)
                        augmented['hand_tiles'] = hand_tiles
                    
                    # 2. 同类型牌的替换（例如万子替换为同数量的筒子）
                    if random.random() < 0.3:
                        augmented['hand_tiles'] = DataAugmenter._suit_replacement(augmented['hand_tiles'])
                
                # 3. 添加噪声（小概率随机修改某些字段）
                if random.random() < 0.1:
                    # 轻微修改状态表示
                    if 'state' in augmented:
                        state = augmented['state'].copy()
                        # 随机调整一些非关键值
                        for key in state:
                            if isinstance(state[key], (int, float)) and random.random() < 0.2:
                                # 小范围扰动
                                state[key] += random.uniform(-0.1, 0.1)
                        augmented['state'] = state
                
                # 添加增强标记
                augmented['augmented'] = True
                augmented_data.append(augmented)
        
        # 合并原始数据和增强数据
        all_data = data + augmented_data
        
        # 打乱顺序
        random.shuffle(all_data)
        
        return all_data
    
    @staticmethod
    def _suit_replacement(hand_tiles: List[str]) -> List[str]:
        """
        花色替换增强
        
        Args:
            hand_tiles: 手牌列表
            
        Returns:
            替换后的手牌列表
        """
        # 定义替换映射
        suit_map = {
            'W': 'T',  # 万 -> 筒
            'T': 'S',  # 筒 -> 索
            'S': 'W'   # 索 -> 万
        }
        
        # 随机选择一个替换方向
        source_suit = random.choice(['W', 'T', 'S'])
        target_suit = suit_map[source_suit]
        
        new_hand = []
        for tile in hand_tiles:
            # 只替换指定花色的牌，跳过字牌
            if len(tile) >= 2 and tile[0] == source_suit and tile[0] in ['W', 'T', 'S']:
                new_tile = target_suit + tile[1:]
                new_hand.append(new_tile)
            else:
                new_hand.append(tile)
        
        return new_hand


class DatasetAnalyzer:
    """
    数据集分析器
    """
    @staticmethod
    def analyze_dataset(data_path: str) -> Dict[str, Any]:
        """
        分析数据集
        
        Args:
            data_path: 数据文件路径
            
        Returns:
            分析结果字典
        """
        # 加载数据
        with open(data_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 基础统计
        total_records = len(data)
        
        # 动作分布
        action_counts = {}
        
        # 玩家分布
        player_counts = {}
        
        # 牌型分布
        tile_type_counts = {}
        
        # 遍历数据
        for record in data:
            # 统计动作
            action = record.get('action', 'unknown')
            action_counts[action] = action_counts.get(action, 0) + 1
            
            # 统计玩家
            player = record.get('player', 'unknown')
            player_counts[player] = player_counts.get(player, 0) + 1
            
            # 统计牌型
            hand_tiles = record.get('hand_tiles', [])
            for tile_str in hand_tiles:
                if len(tile_str) > 0:
                    tile_type = tile_str[0]  # 花色
                    tile_type_counts[tile_type] = tile_type_counts.get(tile_type, 0) + 1
        
        # 计算胜率（如果有游戏结果）
        win_rates = {}
        for record in data:
            result = record.get('game_result', {})
            if 'winner' in result:
                winner = result['winner']
                if winner not in win_rates:
                    win_rates[winner] = {'wins': 0, 'total': 0}
                win_rates[winner]['total'] += 1
                if winner == record.get('player'):
                    win_rates[winner]['wins'] += 1
        
        # 计算胜率百分比
        for player in win_rates:
            if win_rates[player]['total'] > 0:
                win_rates[player]['rate'] = win_rates[player]['wins'] / win_rates[player]['total'] * 100
        
        # 返回分析结果
        analysis = {
            'total_records': total_records,
            'action_distribution': action_counts,
            'player_distribution': player_counts,
            'tile_type_distribution': tile_type_counts,
            'win_rates': win_rates
        }
        
        return analysis
    
    @staticmethod
    def visualize_analysis(analysis: Dict[str, Any], output_file: str) -> None:
        """
        可视化分析结果
        
        Args:
            analysis: 分析结果
            output_file: 输出文件路径
        """
        try:
            import matplotlib.pyplot as plt
            import seaborn as sns
            
            # 设置风格
            sns.set(style="whitegrid")
            
            # 创建图表
            fig, axes = plt.subplots(2, 2, figsize=(15, 12))
            
            # 1. 动作分布
            ax1 = axes[0, 0]
            actions = list(analysis['action_distribution'].keys())
            counts = list(analysis['action_distribution'].values())
            sns.barplot(x=actions, y=counts, ax=ax1)
            ax1.set_title('动作分布')
            ax1.set_xlabel('动作')
            ax1.set_ylabel('数量')
            ax1.tick_params(axis='x', rotation=45)
            
            # 2. 玩家分布
            ax2 = axes[0, 1]
            players = list(analysis['player_distribution'].keys())
            player_counts = list(analysis['player_distribution'].values())
            sns.barplot(x=players, y=player_counts, ax=ax2)
            ax2.set_title('玩家分布')
            ax2.set_xlabel('玩家')
            ax2.set_ylabel('数量')
            
            # 3. 牌型分布
            ax3 = axes[1, 0]
            tile_types = list(analysis['tile_type_distribution'].keys())
            tile_counts = list(analysis['tile_type_distribution'].values())
            sns.barplot(x=tile_types, y=tile_counts, ax=ax3)
            ax3.set_title('牌型分布')
            ax3.set_xlabel('牌型')
            ax3.set_ylabel('数量')
            
            # 4. 胜率分布
            ax4 = axes[1, 1]
            if analysis['win_rates']:
                winners = list(analysis['win_rates'].keys())
                rates = [analysis['win_rates'][w]['rate'] for w in winners]
                sns.barplot(x=winners, y=rates, ax=ax4)
                ax4.set_title('胜率分布')
                ax4.set_xlabel('玩家')
                ax4.set_ylabel('胜率 (%)')
                ax4.set_ylim(0, 100)
            else:
                ax4.text(0.5, 0.5, '无胜率数据', ha='center', va='center', transform=ax4.transAxes)
                ax4.set_title('胜率分布')
            
            # 调整布局
            plt.tight_layout()
            
            # 保存图表
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
            plt.close()
            
            print(f"分析可视化已保存到: {output_file}")
            
        except ImportError as e:
            print(f"无法导入可视化库: {e}")
            print("请安装 matplotlib 和 seaborn 以使用可视化功能")
    
    @staticmethod
    def save_analysis_summary(analysis: Dict[str, Any], output_file: str) -> None:
        """
        保存分析摘要
        
        Args:
            analysis: 分析结果
            output_file: 输出文件路径
        """
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("数据集分析摘要\n")
            f.write("=" * 50 + "\n\n")
            
            # 基础信息
            f.write(f"总记录数: {analysis['total_records']}\n\n")
            
            # 动作分布
            f.write("动作分布:\n")
            for action, count in analysis['action_distribution'].items():
                percentage = count / analysis['total_records'] * 100
                f.write(f"  {action}: {count} ({percentage:.2f}%)\n")
            f.write("\n")
            
            # 玩家分布
            f.write("玩家分布:\n")
            for player, count in analysis['player_distribution'].items():
                percentage = count / analysis['total_records'] * 100
                f.write(f"  {player}: {count} ({percentage:.2f}%)\n")
            f.write("\n")
            
            # 牌型分布
            f.write("牌型分布:\n")
            for tile_type, count in analysis['tile_type_distribution'].items():
                f.write(f"  {tile_type}: {count}\n")
            f.write("\n")
            
            # 胜率分布
            if analysis['win_rates']:
                f.write("胜率分布:\n")
                for player, stats in analysis['win_rates'].items():
                    if 'rate' in stats:
                        f.write(f"  {player}: {stats['rate']:.2f}% ({stats['wins']}/{stats['total']})\n")
            else:
                f.write("无胜率数据\n")


def generate_sample_data(output_path: str, 
                        num_records: int = 100) -> None:
    """
    生成样本数据（用于测试）
    
    Args:
        output_path: 输出路径
        num_records: 记录数量
    """
    # 创建生成器
    generator = MahjongDataGenerator(
        output_dir=os.path.dirname(output_path),
        num_games=num_records // 20,  # 估算每局游戏约20条记录
        seed=42
    )
    
    # 生成数据
    dataset = generator.generate_dataset(expert_only=True)
    
    # 只取指定数量的记录
    dataset = dataset[:num_records]
    
    # 保存数据
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(dataset, f, ensure_ascii=False, indent=2)
    
    print(f"样本数据已生成: {output_path}, 共 {len(dataset)} 条记录")